package com.qen.encrypt.utils;

import cn.hutool.core.codec.Base64;
import cn.hutool.core.util.CharsetUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.Mode;
import cn.hutool.crypto.Padding;
import cn.hutool.crypto.symmetric.AES;
import com.qen.encrypt.exception.EncryptException;
import com.qen.encrypt.properties.AlgorithmEnum;
import lombok.SneakyThrows;

import java.nio.charset.StandardCharsets;
import java.security.SecureRandom;

/**
 * 参数加密工具类:AES-ECB   AES-CBC
 */
public class AESUtils implements ParamEncrypt {
    //    private AES aes;
    private Mode mode;
    private Padding padding;
    private String key;

    public AESUtils(String key, Mode mode, Padding padding) {
        int length = key.length();
        if (16 != length && 24 != length && 32 != length) {
            throw new EncryptException("AES—ECB加密密钥长度不正确");
        }
        this.mode = mode;
        this.padding = padding;
        this.key = key;
//        this.aes = new AES(mode, padding, key.getBytes(StandardCharsets.UTF_8));
    }

    public AESUtils(String key, AlgorithmEnum algorithm) {
        int length = key.length();
        if (16 != length && 24 != length && 32 != length) {
            throw new EncryptException("AES—ECB加密密钥长度不正确");
        }
        Mode mode = algorithm.getMode();
        Padding padding = algorithm.getPadding();
        this.mode = mode;
        this.padding = padding;
        this.key = key;
//        this.aes = new AES(mode, padding, key.getBytes(StandardCharsets.UTF_8));
    }

    @SneakyThrows
    private byte[] initIv() {
        SecureRandom sr = SecureRandom.getInstanceStrong();
        byte[] iv = sr.generateSeed(16);//生成16个字节的随机数
        return iv;
    }

    private byte[] join(byte[] bs1, byte[] bs2) {
        byte[] r = new byte[bs1.length + bs2.length];
        System.arraycopy(bs1, 0, r, 0, bs1.length);
        System.arraycopy(bs2, 0, r, bs1.length, bs2.length);
        return r;
    }

    @Override
    public String pubEncrypt(byte[] data) {
        byte[] iv = initIv();
        AES aes = new AES(mode, padding, key.getBytes(StandardCharsets.UTF_8));
        if (mode.equals(Mode.CBC)) {
            aes.setIv(iv);
            byte[] en = aes.encrypt(data);
            byte[] out = join(iv, en);
            return Base64.encode(out);
        } else {
            return aes.encryptBase64(data);
        }
    }

    @Override
    public String pubEncrypt(String data) {
        return pubEncrypt(StrUtil.bytes(data, CharsetUtil.CHARSET_UTF_8));
    }

    @Override
    public byte[] priDecrypt(byte[] input) {
        AES aes = new AES(mode, padding, key.getBytes(StandardCharsets.UTF_8));
        if (mode.equals(Mode.CBC)) {
            byte[] iv = new byte[16];
            byte[] data = new byte[input.length - 16];
            //iv
            System.arraycopy(input, 0, iv, 0, 16);
            //data
            System.arraycopy(input, 16, data, 0, data.length);
            aes.setIv(iv);
            return aes.decrypt(data);
        } else {
            return aes.decrypt(input);
        }
    }

    @Override
    public String priDecrypt(String data) {
        return new String(priDecrypt(Base64.decode(data)), StandardCharsets.UTF_8);
    }

    @Override
    public String priEncrypt(byte[] data) {
        return pubEncrypt(data);
    }

    @Override
    public String priEncrypt(String data) {
        return pubEncrypt(data);
    }

    @Override
    public byte[] pubDecrypt(byte[] data) {
        return priDecrypt(data);
    }

    @Override
    public String pubDecrypt(String data) {
        return priDecrypt(data);
    }

}
